#include "lidar_2d_undistortion.h"
#include <pcl_ros/transforms.h>

Lidar2DUndistortion::Lidar2DUndistortion(){
    init_params_ = false;
    lidar_delay_time_ = 0.01;   //10ms
    imu_frequency_ = 45;
    imu_max_mums_ = imu_frequency_ * 1.5;
    imu_cache_buffer_ = ImuCircularBuffer(imu_max_mums_);    // IMU 100Hz  预留1.5秒数据

}

// 设置传感器参数数据
void Lidar2DUndistortion::setParams(const SensorParams& sensor_params){
    LOG(INFO) << "Lidar2DUndistortion::setParams" << endl;
    sensor_params_ = sensor_params;
    init_params_ = true;
}

// 添加IMU数据
void Lidar2DUndistortion::addImuData(sensor_msgs::Imu::Ptr data_ptr){
    sensor_msgs::Imu::Ptr new_msg = boost::make_shared<sensor_msgs::Imu>(*data_ptr);
    //建立当前航向的水平坐标系四元数
    new_msg->orientation = imuToCurrentQuaternion(data_ptr->orientation, 3.0);
    imu_cache_buffer_.push_front(new_msg);
}

// 添加里程计数据
void Lidar2DUndistortion::addOdomData(nav_msgs::Odometry::Ptr data_ptr){

}

geometry_msgs::Quaternion Lidar2DUndistortion::imuToCurrentQuaternion(geometry_msgs::Quaternion quat, float yaw_offset){
    tf2::Quaternion current_tf_quat(quat.x, quat.y, quat.z, quat.w);
    //提取航向角
    double roll, pitch, yaw;
    tf2::Matrix3x3(current_tf_quat).getRPY(roll, pitch, yaw);
    tf2::Quaternion quaternion;
    quaternion.setRPY(0.0, 0.0, yaw + yaw_offset);
    return tf2::toMsg(quaternion);
}


Eigen::Quaternionf Lidar2DUndistortion::imuToEigenQuaternion(geometry_msgs::Quaternion quat, float yaw_offset){
    //建立当前航向的水平坐标系四元数
    geometry_msgs::Quaternion horizontal_quat = imuToCurrentQuaternion(quat, yaw_offset);
    return Eigen::Quaternionf(horizontal_quat.w, horizontal_quat.x, horizontal_quat.y, horizontal_quat.z);
}


bool Lidar2DUndistortion::getLaserPose(int point_index, ros::Time& timestamp, Eigen::Quaternionf& quat_out){
    static int index_front = 0;
    static int index_back = 0;
    static ros::Time timestamp_front;
    static ros::Time timestamp_back;
    static Eigen::Quaternionf quat_front, quat_back;

    static bool index_updated_flag = false;
    static bool predict_orientation_flag = false;

    cout << "[getLaserPose] index point_index: " << point_index << " size: " << imu_cache_buffer_.size() << endl;

    if (point_index == 0)
    {
        predict_orientation_flag = false;
        int i = 0;
        while (timestamp < imu_cache_buffer_[i]->header.stamp && i < imu_cache_buffer_.size())  
        {
            i++;
        }
        index_front = i - 1;
        index_back = i;
        index_updated_flag = true;
    }
    else
    {
        cout << "index while: " << index_front << " " << index_back << endl;
        while (predict_orientation_flag == false
                && timestamp > imu_cache_buffer_[index_front]->header.stamp
                && timestamp > imu_cache_buffer_[index_back]->header.stamp)
        {
            
            index_front--;
            index_back--;

            if (index_front < 0)
            {
                cout << "index predict_orientation_flag: " << index_front << " " << index_back << endl;
                //use prediction
                predict_orientation_flag = true;
                index_front++;
                index_back++;
            }

            index_updated_flag = true;
        }
    }

    if (index_updated_flag == true)
    {
        cout << imu_cache_buffer_.size() << ", index updated: " << index_front << " " << index_back << endl;
        timestamp_front = imu_cache_buffer_[index_front]->header.stamp;
        timestamp_back = imu_cache_buffer_[index_back]->header.stamp;
        quat_front = Eigen::Quaternionf(imu_cache_buffer_[index_front]->orientation.w, 
            imu_cache_buffer_[index_front]->orientation.x, 
            imu_cache_buffer_[index_front]->orientation.y, 
            imu_cache_buffer_[index_front]->orientation.z);
        quat_back = Eigen::Quaternionf(imu_cache_buffer_[index_back]->orientation.w, 
            imu_cache_buffer_[index_back]->orientation.x, 
            imu_cache_buffer_[index_back]->orientation.y, 
            imu_cache_buffer_[index_back]->orientation.z);

        index_updated_flag = false;
    }

    float alpha = (float)(timestamp.toNSec() - timestamp_back.toNSec()) / (timestamp_front.toNSec() - timestamp_back.toNSec());

    if (alpha < 0)
    {
        return false;
    }

    // 球面线性插值
    // Slerp.
    quat_out = quat_back.slerp(alpha, quat_front);
    return true;
}

// 获取去除畸变数据
pcl::PointCloud<pcl::PointXYZ>::Ptr Lidar2DUndistortion::getLidarData(sensor_msgs::LaserScan::Ptr data_input){
    if(imu_cache_buffer_.size() < imu_max_mums_){
        ROS_INFO_STREAM("need more imu data. imu_cache_buffer_ size is " 
            << imu_cache_buffer_.size());
        return nullptr;
    }    

    // 激光点的个数
    int length = data_input->ranges.size();
    ros::Duration frame_duration = ros::Duration(data_input->scan_time);
    ros::Time lidar_timebase = data_input->header.stamp - ros::Duration(lidar_delay_time_);

    // 将LaserScan类型的数据转化为PCL点云格式
    pcl::PointCloud<pcl::PointXYZ>::Ptr cloud_ptr(new pcl::PointCloud<pcl::PointXYZ>());
    pcl::PointCloud<pcl::PointXYZ>::Ptr pointcloud_raw_pcl(new pcl::PointCloud<pcl::PointXYZ>());
    for(size_t i=0; i<data_input->ranges.size(); ++i){
        float range = data_input->ranges[i];
        float angle = data_input->angle_min + i * data_input->angle_increment;
        pcl::PointXYZ point;
        point.x = range * cos(angle);
        point.y = range * sin(angle);
        point.z = 0.0f;
        pointcloud_raw_pcl->points.emplace_back(point);
    }

    if(data_input->scan_time = 0){
        ROS_WARN_STREAM("scan_time is 0");
        return pointcloud_raw_pcl;
    }

    pcl::PointCloud<pcl::PointXYZ>::Ptr out_ptr(new pcl::PointCloud<pcl::PointXYZ>());
    if (imu_cache_buffer_.size() > frame_duration.toSec() * imu_max_mums_){
        auto start = system_clock::now();
        // pcl::PointCloud<pcl::PointXYZI> pointcloud_pcl;
        // pcl::PointXYZI point_xyzi;
        pcl::PointXYZ point;
        // ros::Time current_output_timestamp = imu_cache_buffer_[0]->header.stamp;
        Eigen::Quaternionf current_quat = imuToEigenQuaternion(imu_cache_buffer_[0]->orientation, 0.0);
        
        for (int i = 0; i < length; i++){
            // 获取当前点的时间戳
            ros::Time point_timestamp = lidar_timebase + ros::Duration(data_input->time_increment * i);  
            // 获取当前点位姿
            pcl::PointXYZ current_point = pointcloud_raw_pcl->points[i];
            Eigen::Vector3f point_in(current_point.x, current_point.y, 0.0);

           // 构造输入输出位姿
           Eigen::Quaternionf point_quat;
            // 如果成功获取当前扫描点的姿态
            if (getLaserPose(i, point_timestamp, point_quat) == true){
                Eigen::Quaternionf delta_quat = current_quat.inverse() * point_quat;

                Eigen::Vector3f point_out = delta_quat * point_in;
                point.x = point_out(0);
                point.y = point_out(1);
                point.z = 0.0;
                // point_xyzi.intensity = current_point.intensity;
                out_ptr->points.push_back(point);
            }  
            else{
                break;
            } 
        }
    } 
    return out_ptr;
}

// 点云坐标转换
pcl::PointCloud<pcl::PointXYZ>::Ptr Lidar2DUndistortion::transformPCL(pcl::PointCloud<pcl::PointXYZ>::Ptr data_input){
    // 获取变换信息
    // cout << "sensor_params_.transform.yaw: " << sensor_params_.transform.yaw << endl;
    Eigen::Translation3d translation(sensor_params_.transform.x, sensor_params_.transform.y, sensor_params_.transform.z);
    Eigen::Matrix3d R = 
        Eigen::AngleAxisd(sensor_params_.transform.yaw, Eigen::Vector3d::UnitZ()).toRotationMatrix() *
        Eigen::AngleAxisd(sensor_params_.transform.pitch, Eigen::Vector3d::UnitY()).toRotationMatrix() *
        Eigen::AngleAxisd(sensor_params_.transform.roll, Eigen::Vector3d::UnitX()).toRotationMatrix();
    Eigen::Matrix4d transformMatrix = Eigen::Matrix4d::Identity();
    transformMatrix.block<3, 3>(0, 0) = R;
    transformMatrix.block<3, 1>(0, 3) = translation.vector();
    Eigen::Matrix4d inverseMatrix = transformMatrix.inverse();
    
    // 点云数据转换到base_link坐标系
    pcl::PointCloud<pcl::PointXYZ>::Ptr data_output(new pcl::PointCloud<pcl::PointXYZ>());
    pcl::transformPointCloud(*data_input, *data_output, transformMatrix);
    return data_output;
}
